import torch
import numpy as np
from sklearn.neighbors import KNeighborsClassifier
from sklearn.metrics import accuracy_score
import matplotlib.pyplot as plt

# 生成一个简单的数据集 (2个特征和2个分类)
# X为输入特征，y为标签
X = np.array([[1, 2], [2, 3], [3, 4], [5, 7], [6, 8], [7, 9], [8, 10], [3, 6], [4, 5], [6, 4]])
y = np.array([0, 0, 0, 1, 1, 1, 1, 0, 0, 1])

# 数据转换为 PyTorch 张量
X_tensor = torch.tensor(X, dtype=torch.float32)
y_tensor = torch.tensor(y, dtype=torch.long)

# 打印数据
print("Features:")
print(X_tensor)
print("Labels:")
print(y_tensor)

# 使用 sklearn KNN 分类器，调整邻居数量为 5
knn = KNeighborsClassifier(n_neighbors=5)
knn.fit(X, y)

# 预测
y_pred = knn.predict(X)

# 计算准确率
accuracy = accuracy_score(y, y_pred)
print(f"Accuracy: {accuracy * 100:.2f}%")

# 可视化数据
plt.figure(figsize=(6, 4))
plt.scatter(X[:, 0], X[:, 1], c=y, cmap='bwr', marker='o', edgecolor='k', s=100)
plt.title("KNN Classification Example")
plt.xlabel("Feature 1")
plt.ylabel("Feature 2")
plt.show()

# 测试：给定新的输入数据进行预测
test_data = np.array([[5, 6], [2, 3]])
test_prediction = knn.predict(test_data)

print(f"Predictions for test data {test_data} are {test_prediction}")
